#include "QLearningAlgorithm.h"
#include "MapGen.h"

QLearningAlgorithm::QLearningAlgorithm(MapGen *mapgen) :
    m_mapGen(mapgen)
{
    iniData();
    refreshHashQ();
}

QLearningAlgorithm::~QLearningAlgorithm()
{

}

void QLearningAlgorithm::refreshHashQ()
{
    int num = 0;
    while(num < TRYNUM)
    {
        int player_x = 0;
        int player_y = 0;
        m_mapGen->genPlayer(player_x, player_y);

        //random walk and refresh Q table
        while(m_mapGen->m_map[player_x][player_y] != '0')
        {
            int playerstate_i = ROOM_WIDTH*player_x + player_y;

            //choose a direction, find next state index
            int direction = 0;
            int nextstate_i = 0;
            vecstates[playerstate_i].actRandom(direction, nextstate_i);

            int curQ = hashQ.find(playerstate_i).value()[direction];
            int index = 0;
            hashQ.find(playerstate_i).value()[direction] = curQ + LRATE * (hashR.find(playerstate_i).value()[direction] +
                                                                           ALPHA * maxQVec(hashQ.find(nextstate_i).value(), index) - curQ) ;

            player_x = vecstates[nextstate_i].m_x;
            player_y = vecstates[nextstate_i].m_y;
        }
        num++;
    }
}

void QLearningAlgorithm::iniQHashWithFloat(QHash<int, QVector<float> > &hash, float f)
{
    for(int i = 0; i < ROOM_WIDTH; i++)
    {
        for(int j = 0; j < ROOM_HEIGHT; j++)
        {
            int curstatei = i * ROOM_WIDTH + j;
            QVector<float> v;
            v.push_back(f);v.push_back(f);v.push_back(f);v.push_back(f);
            hash.insert(curstatei, v);
        }
    }
}

void QLearningAlgorithm::iniData()
{
    iniQHashWithFloat(hashQ, 0.0f);
    iniQHashWithFloat(hashR, -100.0f);

    for(int i = 0; i < ROOM_WIDTH; i++)
    {
        for(int j = 0; j < ROOM_HEIGHT; j++)
        {
            State state;
            setState(state, i, j, m_mapGen->m_map);
            vecstates.push_back(state);
        }
    }

    int statenum = vecstates.size();
    for(int statei = 0; statei < statenum; statei++)
    {
        QMap<int, int> *action = &vecstates[statei].m_transfer;
        for(auto it = action->begin(); it != action->end(); ++it)
        {
            int dir = it.key();
            int nexti = it.value();

            float qvalue = 0.f;
            switch(vecstates[nexti].m_tile)
            {
            case '#': qvalue = R_OUT; break;
            case 'M': qvalue = R_MONS; break;
            case '0': qvalue = R_EXIT; break;
            case 'T': qvalue = R_TREA; break;
            default: qvalue = R_STEP; break;
            }
            hashR.find(statei).value()[dir] = qvalue;
        }
    }
}

void QLearningAlgorithm::setState(State &state, int x, int y, char m_map[ROOM_WIDTH][ROOM_HEIGHT])
{
    state.iniData(x, y);
    state.m_tile = m_map[x][y];

    bool coli_up = (x - 1) < 0 || m_map[x-1][y] == '#';
    if(!coli_up)
    {
        int iup = ROOM_WIDTH * (x - 1) + y;
        state.m_transfer.insert(0, iup);
    }

    bool coli_do = (x + 1) >= ROOM_WIDTH || m_map[x+1][y] == '#';
    if(!coli_do)
    {
        int ido = ROOM_WIDTH * (x + 1) + y;
        state.m_transfer.insert(1, ido);
    }

    bool coli_le = (y - 1) < 0 || m_map[x][y-1] == '#';
    if(!coli_le)
    {
        int ile = ROOM_WIDTH * x + y - 1;
        state.m_transfer.insert(2, ile);
    }

    bool coli_ri = (y + 1) >= ROOM_HEIGHT || m_map[x][y+1] == '#';
    if(!coli_ri)
    {
        int iri = ROOM_WIDTH * x + y + 1;
        state.m_transfer.insert(3, iri);
    }
}

int QLearningAlgorithm::maxQDir(int statei)
{
    int index = 0;
    float maxvec = -999.f;
    QVector<int> possilbedir;
    for(auto it = vecstates[statei].m_transfer.begin(); it != vecstates[statei].m_transfer.end(); ++it)
    {
        possilbedir.push_back(it.key());
    }
    for(int i = 0; i < possilbedir.size(); i++)
    {
        int dir = possilbedir[i];
        float qvalue = hashQ.find(statei).value()[dir];
        if(qvalue > maxvec)
        {
            maxvec = qvalue;
            index = dir;
        }
    }
    return index;
}

float maxQVec(QVector<float> &vec, int &index)
{
    int num = vec.size();
    index = 0;
    float maxvec = -999.f;
    for(int i = 0; i < num; i++)
    {
        if(vec[i] > maxvec)
        {
            maxvec = vec[i];
            index = i;
        }
    }
    return maxvec;
}
